import tensorflow as tf


def get_mask_3d(seq_lens, max_len, dtype=tf.float32):
    """Mask for feature map of shape [Batch_size, max_time, channel]"""
    mask = tf.cast(tf.sequence_mask(seq_lens, max_len), dtype=dtype)
    mask = tf.expand_dims(mask, -1)
    return mask


def get_mask_4d(seq_lens, max_len, dtype=tf.float32):
    """
    Mask for feature map of shape [batch_size, max_time, dim(freq), channel]
    """
    mask = tf.cast(tf.sequence_mask(seq_lens, max_len), dtype=dtype)
    mask = tf.expand_dims(tf.expand_dims(mask, -1), -1)
    return mask


def conv_block(inputs, filters=64, kernel_size=(3, 3), h_id=None, g_id=None,
               padding='same'):
    outputs = tf.layers.conv2d(inputs=inputs, filters=filters * 2,
                               kernel_size=kernel_size, padding=padding)
    a, b = tf.split(outputs, 2, -1)
    if h_id:
        a = a + h_id
    if g_id:
        b = b + g_id
    outputs = tf.nn.tanh(a) * tf.nn.sigmoid(b)
    return outputs


def avg_pooling2d(inputs,
                  pool_size=(2, 2),
                  strides=(2, 2),
                  padding='valid',
                  name=None):
    h = tf.layers.average_pooling2d(inputs=inputs,
                                    pool_size=pool_size,
                                    strides=strides,
                                    padding=padding,
                                    name=name)
    return h


def deconvolution(inputs):
    return tf.layers.conv2d_transpose(inputs, filters=64, kernel_size=(2, 2),
                                      strides=(2, 2), padding='same')
